from redbaron import RedBaron
import argparse
from typing import Any, Optional
from dataclasses import dataclass
import abc
from typing import Iterable

from collections import defaultdict
from loguru import logger


@dataclass      
class BasicVarInfo:
    path: list[int|str] = None
    type: abc.ABCMeta = None
    value: Any = None
    nargs: str|int = 1 
    help: Optional[str] = None
    var: list[str] = None
    
    def __post_init__(self):
        self.value: Optional[self.type]
        self.name = self.var[-1]

# class ParserRegister:
#     def __init__(self, 
#                  task: str,
#                  config_path: str
#                  ):
        
#         self.args_parse_helper = ArgsParseHelper(task, config_path)
#         var_and_info = self.args_parse_helper.get_var_and_info()
        
#         # 必须的传参
#         parser = argparse.ArgumentParser()
        
#         for (var_hash, var_info) in var_and_info:
#             _var, _type, _value, _nargs, _help = (
#                 var_info.name,
#                 var_info.type,
#                 var_info.value,
#                 var_info.nargs,
#                 var_info.help
#                 )
            
#             parser.add_argument(
#                 f'--{var_hash}', 
#                 type=_type,
#                 default=_value,
#                 nargs=_nargs,
#                 help=_help
#             )
            
#         self.args = parser.parse_args()
        
            
#     def run(self):
#         self.args_parse_helper.run(self.args)
    
        
# class ArgsParseHelper:
#     task2var2path: dict[str, dict[str, BasicVarInfo]] = {}
#     task2redbaron: dict[str, RedBaron] = {}
    
#     task2config = dict(
#             train='/mnt/user/linzhixin/heatmap/config/train/cfg_qwen2_7b_sft.py',
#             trainer='/mnt/user/linzhixin/heatmap/config/train/cfg_qwen2_7b_sft.py',
#             misc='/mnt/user/linzhixin/heatmap/config/misc/cfg_convert_qwen2_mgtcore_to_hf.py'
#         )
    
#     def __init__(self, 
#                  task: str,
#                  config_path: Optional[str]):
#         self.curr_task = task 
#         self.curr_config_path = config_path

#         self.task2config[task] = config_path
#         self.init_var2path()
        
#     def run(self,
#             args,
#             save_path: Optional[str] = None):
#         """
#             1.读取原config的redbaron
#             2.按task找对应的解析路径
#             3.按路径和对应var_hash修改redbaron的值
#             4.dump redbaron
#         """
        
#         def path4redbaron(basic_info: BasicVarInfo, 
#                           new_value: Any, 
#                           red):
            
#             if new_value and len(new_value) == 1:
#                 new_value = new_value[0]
                
#             if new_value and not isinstance(new_value, basic_info.type):
#                 raise TypeError(f'类型错误. new_value:{new_value}, type:{type(new_value)}, basic_info.type:{basic_info.type}')
            
#             curr = red
#             # print(path, new_value)
#             # red.help()
#             for step in basic_info.path:
#                 if isinstance(step, int):
#                     curr = curr[step]
#                 elif isinstance(step, str):
#                     if step == 'value':  
#                         curr = curr.value 
#                     else:
#                         raise ValueError(f'Step wrong. Got {step}.')
#             if basic_info.type == str:
#                 new_value = f'"{new_value}"'
#             curr.value = new_value
#             return red 
        
#         red = self.task2redbaron[self.curr_task]
#         var_and_basic_info = self.task2var2path[self.curr_task]
        
#         for var_hash, basic_info in var_and_basic_info.items(): 
#             new_value = getattr(args, var_hash, None)
#             if basic_info.nargs != 1:
#                 new_value = basic_info.type(new_value)
#             if not new_value:
#                 continue
#             red = path4redbaron(basic_info, new_value, red)
        
#         save_path = save_path if save_path else self.task2config[self.curr_task]
        
#         with open(save_path, "w") as source_code:
#             source_code.write(red.dumps())
            
#         print(f'新config存储位置：{save_path}')
        
#     def get_var_and_info(self, task: Optional[str]) -> list[tuple[str, BasicVarInfo]]:
#         if not task:
#             task = self.curr_task
             
#         if task not in self.task2var2path:
#             raise ValueError(f'Task {task} is not in default task list.')
        
#         return [(key,value) for key, value in self.task2var2path[task].items()] 
        

#     def init_var2path(self, 
#                       task: Optional[str]=None, 
#                       config_path: Optional[str]=None):
#         def dfs(red, path=[], var=[]):
#             if hasattr(red, 'target'):
#                 # print('有target')
#                 # print('路径:', red.target.value)
#                 path.append('value')
#                 var.append(red.target.value)
#                 dfs(red.value, path=path, var=var)
#             elif red.find('NameNode') != -1:
#                 if 'value' in red.__dict__.keys():
#                     path.append('value')
#                     # print('value深入')
#                     dfs(red.value, path=path, var=var)
#                 elif isinstance(red, Iterable): 
#                     # print('可迭代')
#                     for index, i in enumerate(red):
#                         dfs(i, path=path+[index], var=var[:])
#             else:
#                 default_value = red.value
#                 if isinstance(default_value, Iterable):
#                     self.task2var2path[task]['-'.join(var)] = BasicVarInfo(
#                         path=path,
#                         type=type(default_value),
#                         default=default_value,
#                         nargs='*',
#                         var=var
#                     ),
#                 else:
#                     self.task2var2path[task]['-'.join(var)] = BasicVarInfo(
#                         path=path,
#                         type=type(default_value),
#                         default=default_value,
#                         var=var
#                     ),
                
#         if not task:
#             task = self.curr_task
                     
#         if not config_path:
#             if task in self.task2config:
#                 config_path = self.task2config[task]
#             else:
#                 raise ValueError(f'Task {task} is not in default task list.')
        
#         with open(config_path, "r") as source_code:
#             red = RedBaron(source_code.read())
            
#         self.task2redbaron[task] = red
#         _red = red.__getattribute__(task)
#         self.task2var2path[task] = {}    
#         dfs(_red)

def load_redbaron(path: str) -> RedBaron:
    with open(path, "r") as source_code:
        red = RedBaron(source_code.read())
    return red

def save_redbaron(red: RedBaron, 
                  save_path: str) -> None:
    with open(save_path, "w") as source_code:
        source_code.write(red.dumps())
        
def redbaron_config_parser(red: RedBaron) -> dict[str, BasicVarInfo]:
    ret = {}
    def dfs(red, path=[], var=[]):
        if hasattr(red, 'target'):
            # print('有target')
            # print('路径:', red.target.value)
            path.append('value')
            var.append(red.target.value)
            dfs(red.value, path=path, var=var)
        elif red.find('NameNode') != -1:
            if 'value' in red.__dict__.keys():
                path.append('value')
                # print('value深入')
                dfs(red.value, path=path, var=var)
            elif isinstance(red, Iterable): 
                # print('可迭代')
                for index, i in enumerate(red):
                    dfs(i, path=path+[index], var=var[:])
        else:
            default_value = red.value
            if isinstance(default_value, Iterable):
                ret[task]['-'.join(var)] = BasicVarInfo(
                    path=path,
                    type=type(default_value),
                    default=default_value,
                    nargs='*',
                    var=var
                ),
            else:
                ret[task]['-'.join(var)] = BasicVarInfo(
                    path=path,
                    type=type(default_value),
                    default=default_value,
                    var=var
                ),
                
    return ret
    
        
        

class ParserRegisterDataHelper:
    """
            用于帮忙管理一些默认数据，也可更新数据
    """
    
    task2config = dict(
            train='/mnt/user/linzhixin/heatmap/config/train/cfg_qwen2_7b_sft.py',
            trainer='/mnt/user/linzhixin/heatmap/config/train/cfg_qwen2_7b_sft.py',
            misc='/mnt/user/linzhixin/heatmap/config/misc/cfg_convert_qwen2_mgtcore_to_hf.py'
        )
    
    default_task: Optional[str] = None 
    
    
    def regist(self, 
               task: str, 
               config_path: str,
               is_default_task: bool=False,):
        if task in self.task2config:
            logger.info(f'Task {task} have default config_path: {self.task2config[task]}')
        else:
            logger.info(f'No default task called {task}.')
        
        logger.info(f'set new task and config_path: {task} | {config_path}')
        self.task2config[task] = config_path
        
        if is_default_task:
            self.default_task = task
    
    def get_config_path(self, task: Optional[str]=None):
        if not task:
            if self.default_task: 
                raise ValueError(f'The task is empty and no default_task.')
            else:
                task = self.default_task 
        
        if task not in self.task2config:
            raise ValueError('The task not in the existing_config. '
                             'Regist task first or select the default task.')
        
        return self.task2config[task]




class ParserRegister:
    def __init__(self, 
                 task: str, 
                 config_path: Optional[str]):
        # 给一个task，根据task可以找默认的config_path，也可以自定义config_path。
        self.task = task
        self.data_helper = ParserRegisterDataHelper()
        if config_path:
            self.data_helper.regist(task, config_path)
        
        # 用redbaron加载config_path。
        self.redbaron = load_redbaron(self.data_helper.get_config_path())
        
        # 解析并注册所有config变量，记载redbaron修改路径、变量路径、默认值、类型，存到BasicVarInfo里。
        self.parsed_redbaron = redbaron_config_parser(self.redbaron)
        
        
        # 用argparse注册所有redbaron解析出来的修改路径，并设置传参类型必须为string，到时候都用string修改具体的值。
        
        
        # parser解析所有传参，并根据BasicVarInfo修改default，再根据更新后的BasicVarInfo们更新redbaron。
        
        
        # dump redbaron，结束。
        
        ...
        
        
        

parser_register = ParserRegister

if __name__ == '__main__':
    import sys 
    task = sys.argv[1]
    config_path = sys.argv[2]
    
    parser = parser_register(task, config_path)
    parser.run()
